# import os
# os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from transformers import BertTokenizer, BartForConditionalGeneration,Text2TextGenerationPipeline


tokenizer = BertTokenizer.from_pretrained("model/bart-base-chinese")
model = BartForConditionalGeneration.from_pretrained("model/bart-base-chinese")
text2text_generator = Text2TextGenerationPipeline(model, tokenizer)
text2text_generator("北京是[MASK]的首都", max_length=50, do_sample=False)

# 输入文本和摘要生成
input_text ="""
案情简介：李女士报警称此前在社交平台结识一名网友，对
方自称是现役军人，一番“深情”输出将李女士迷惑。之后对方以服役不方便使用手机为由托李女士帮忙在一网站购买黄金并提供了登录账号及密码(网站为骗子的钓鱼网站)。
李女士欣然答应并向该网站转账30余万元用于购买黄金。随后对方开始诱导李女士一起参与投资。李女士心动了，在对方的“甜蜜攻势”下多次向一银行账户转款40余万用于投资，直到多次转账后并未收到黄金，李女士才意识到被骗了，共计被诈骗70余万元。
案件启示：犯罪分子借助婚恋网站或交友软件寻找目标，假借与受害人谈恋爱为名添加对方微信，取得受害人信任后，以投资赚钱为诱饵诱使受害人在虚假平台投资，最终转移资金、卷款消失，导致受害人“人财两空”，感情、金钱双双损失惨重。
网络交友时，要提高警惕，不要轻易相信陌生人，尤其是那些在短时间内就表现出过度热情的人。对“网恋对象”说的高额回报，一定要不听不信不转账。"""
input_ids = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True).input_ids
outputs = model.generate(input_ids, max_length=150, min_length=50, num_beams=4, length_penalty=2.0)
output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(output_text)